package com.thinkaurelius.titan.testutil.gen;

import java.util.HashMap;
import java.util.Map;

import com.google.common.base.Preconditions;
import com.thinkaurelius.titan.core.TitanGraph;
import com.thinkaurelius.titan.core.TitanKey;
import com.thinkaurelius.titan.core.TitanTransaction;
import com.tinkerpop.blueprints.Edge;
import com.tinkerpop.blueprints.Vertex;

public class Schema {

    public static final String VERTEX_KEY_PREFIX = "vp_";
    public static final String EDGE_KEY_PREFIX = "ep_";
    public static final String LABEL_PREFIX = "el_";
    public static final String UID_PROP = "uid";

    public static long SUPERNODE_UID = 0L;
    private static final int SUPERNODE_INDEX = 0;

    private final int edgeCount;
    private final int vertexCount;
    private final int maxEdgePropVal;
    private final int maxVertexPropVal;
    /*
     * edgeCount must have type int instead of long because
     * DistributionGenerator expects int. It's probably not a great idea to go
     * over 4B per label in memory anyway.
     */
    private final int vertexPropKeys;
    private final int edgePropKeys;
    private final int edgeLabels;
    private final String[] vertexPropNames;
    private final String[] edgePropNames;
    private final String[] edgeLabelNames;
    private final Map<String, String> labelPkeys;

    public static class Builder {

        private int maxVertexPropVal = 100;
        private int maxEdgePropVal = 100;
        private int vertexPropKeys = 20;
        private int edgePropKeys = 10;
        private int edgeLabels = 3;
        private int vertexCount = -1;
        private int edgeCount = -1;

        /**
         * Set the maximum value of vertex properties. This is an exclusive
         * limit. The minimum is always 0.
         *
         * @param m maximum vertex property value, exclusive
         * @return self
         */
        public Builder setMaxVertexPropVal(int m) {
            maxVertexPropVal = m;
            return this;
        }

        /**
         * Set the maximum value of edge properties. This is an exclusive limit.
         * The minimum is always 0.
         *
         * @param m maximum edge property value, exclusive
         * @return self
         */
        public Builder setMaxEdgePropVal(int m) {
            maxEdgePropVal = m;
            return this;
        }

        /**
         * Set the total number of distinct property keys to use for vertex
         * properties.
         *
         * @param vertexPropKeys number of property keys
         * @return self
         */
        public Builder setVertexPropKeys(int vertexPropKeys) {
            this.vertexPropKeys = vertexPropKeys;
            return this;
        }

        /**
         * Set the total number of distinct property keys to use for edge
         * properties.
         *
         * @param edgePropKeys number of property keys
         * @return self
         */
        public Builder setEdgePropKeys(int edgePropKeys) {
            this.edgePropKeys = edgePropKeys;
            return this;
        }

        /**
         * Set the total number of edge labels to create.
         *
         * @param edgeLabels number of edge labels
         * @return self
         */
        public Builder setEdgeLabels(int edgeLabels) {
            this.edgeLabels = edgeLabels;
            return this;
        }

        /**
         * Set the number of vertices to create.
         *
         * @param vertexCount global vertex total
         * @return self
         */
        public Builder setVertexCount(int vertexCount) {
            this.vertexCount = vertexCount;
            Preconditions.checkArgument(0 <= this.vertexCount);
            return this;
        }

        /**
         * Set the number of edges to create for each edge label.
         *
         * @param edgeCount global edge total for each label
         * @return self
         */
        public Builder setEdgeCount(int edgeCount) {
            this.edgeCount = edgeCount;
            Preconditions.checkArgument(0 <= this.edgeCount);
            return this;
        }

        public Builder(int vertexCount, int edgeCount) {
            setVertexCount(vertexCount);
            setEdgeCount(edgeCount);
        }

        /**
         * Construct a {@link GraphGenerator} with this {@code Builder}'s
         * settings.
         *
         * @return a new GraphGenerator
         */
        public Schema build() {
            return new Schema(maxEdgePropVal, maxVertexPropVal, vertexCount, edgeCount, vertexPropKeys, edgePropKeys, edgeLabels);
        }
    }

    public final String getVertexPropertyName(int i) {
        return vertexPropNames[i];
    }

    public final String getEdgePropertyName(int i) {
        return edgePropNames[i];
    }

    public final String getEdgeLabelName(int i) {
        return edgeLabelNames[i];
    }

    public final String getSortKeyForLabel(String l) {
        return l.replace("el_", "ep_");
    }

    public final int getVertexPropKeys() {
        return vertexPropKeys;
    }

    public final int getEdgePropKeys() {
        return edgePropKeys;
    }

    public final int getMaxEdgePropVal() {
        return maxEdgePropVal;
    }

    public final int getMaxVertexPropVal() {
        return maxVertexPropVal;
    }

    public final int getEdgeLabels() {
        return edgeLabels;
    }

    public final long getSupernodeUid() {
        return SUPERNODE_UID;
    }

    public final String getSupernodeOutLabel() {
        return getEdgeLabelName(SUPERNODE_INDEX);
    }

    public final long getMaxUid() {
        return vertexCount;
    }

    public final int getVertexCount() {
        return vertexCount;
    }

    public final int getEdgeCount() {
        return edgeCount;
    }

    private Schema(int maxEdgePropVal, int maxVertexPropVal, int vertexCount,
                   int edgeCount, int vertexPropKeys, int edgePropKeys, int edgeLabels) {
        this.maxEdgePropVal = maxEdgePropVal;
        this.maxVertexPropVal = maxVertexPropVal;
        this.vertexCount = vertexCount;
        this.edgeCount = edgeCount;
        this.vertexPropKeys = vertexPropKeys;
        this.edgePropKeys = edgePropKeys;
        this.edgeLabels = edgeLabels;

        this.vertexPropNames = generateNames(VERTEX_KEY_PREFIX, this.vertexPropKeys);
        this.edgePropNames = generateNames(EDGE_KEY_PREFIX, this.edgePropKeys);
        this.edgeLabelNames = generateNames(LABEL_PREFIX, this.edgeLabels);

        Preconditions.checkArgument(this.edgeLabels <= this.edgePropKeys);

        this.labelPkeys = new HashMap<String, String>(this.edgeLabels);
        for (int i = 0; i < this.edgeLabels; i++) {
            labelPkeys.put(edgeLabelNames[i], edgePropNames[i]);
        }
    }


    public void makeTypes(TitanGraph g) {
        Preconditions.checkArgument(edgeLabels <= edgePropKeys);

        TitanTransaction tx = g.newTransaction();
        for (int i = 0; i < vertexPropKeys; i++) {
            tx.makeKey(getVertexPropertyName(i)).dataType(Integer.class).indexed(Vertex.class).single().make();
        }
        for (int i = 0; i < edgePropKeys; i++) {
            tx.makeKey(getEdgePropertyName(i)).dataType(Integer.class).indexed(Edge.class).single().make();
        }
        for (int i = 0; i < edgeLabels; i++) {
            String labelName = getEdgeLabelName(i);
            String pkName = getSortKeyForLabel(labelName);
            TitanKey pk = tx.getPropertyKey(pkName);
            tx.makeLabel(getEdgeLabelName(i)).sortKey(pk).make();
        }

        tx.makeKey(UID_PROP).dataType(Long.class).indexed(Vertex.class).single().unique().make();
        tx.commit();
    }

    private String[] generateNames(String prefix, int count) {
        String[] result = new String[count];
        StringBuilder sb = new StringBuilder(8);
        sb.append(prefix);
        for (int i = 0; i < count; i++) {
            sb.append(i);
            result[i] = sb.toString();
            sb.setLength(prefix.length());
        }
        return result;
    }
}
